import torch
import torch.nn as nn


# input

wb = torch.tensor([[10.0]])
f = torch.tensor([[2],[2.6],[4],[5]])
f = f.reshape(-1,1)
# label
# TODO: label应该是使用max-min fairness算法得到的标准答案



# Model
class Max_Min_Model(nn.Module):

    def __init__(self):
        super(Max_Min_Model,self).__init__()
        self.model = Model()
        self.output = nn.Linear(1, 4)

    def forward(self,wd,f):  # wd: 1x1, f: 4*1
        zero = torch.tensor([[0.0]])
        wd = torch.cat((zero,wd), 0)  # wd:2*1
        out1 = self.model(wd,f)  # 1*1
        wd_1=wd[1].view(1,-1)
        # wd = torch.cat((out1,wd[1]), 0)  # TODO:tensor维数还有问题，切片后维数下降，需要回到2维
        wd = torch.cat((out1, wd_1), 0)
        out2 = self.model(wd,f)
        # wd = torch.cat((out2, wd[1]), 0)
        wd = torch.cat((out2, wd_1), 0)
        out3 = self.model(wd,f)

        wd = torch.cat((out3, wd_1), 0)
        out4 = self.model(wd, f)
        out = torch.cat((out1,out2,out3,out4), 0)
        return out

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.first = First_layer()
        self.second = Second_layer()
        self.third = nn.Linear(4,1)

    def forward(self,wd,f):
        first = self.first(wd)  # 4*1
        second = self.second(first,f) # 4*1
        output = self.third(second)  # 1*1
        return output


class First_layer(nn.Module):
    def __init__(self):
        super(First_layer,self).__init__()
        self.first = nn.Linear(2,4)

    def forward(self,wd):
        wd=wd.reshape(-1,2)
        return self.first(wd)

class Second_layer(nn.Module):
    def __init__(self):
        super(Second_layer,self).__init__()
        self.s1 = nn.Linear(2, 1)
        self.s2 = nn.Linear(2, 1)
        self.s3 = nn.Linear(2, 1)
        self.s4 = nn.Linear(2, 1)

    def forward(self,first,flows):
        first = first.reshape(-1,1)
        t_tuple = torch.split(first,1)
        f_tuple = torch.split(flows,1)
        f1 = torch.cat((t_tuple[0], f_tuple[0]), 0)
        f2 = torch.cat((t_tuple[1], f_tuple[1]), 0)
        f3 = torch.cat((t_tuple[2], f_tuple[2]), 0)
        f4 = torch.cat((t_tuple[3], f_tuple[3]), 0)
        out1 = self.s1(f1.reshape(1,-1))
        out2 = self.s2(f2.reshape(1,-1))
        out3 = self.s3(f3.reshape(1,-1))
        out4 = self.s4(f4.reshape(1,-1))
        return torch.cat((out1,out2,out3,out4),1)


net = Max_Min_Model()
print(net(wb,f))

# epochs = 5
# lr=0.01
# optimizer = torch.optim.SGD(net.parameters(), lr=lr)
# criterion = nn.MSELoss
#
# # train
# for epoch in range(epochs):
#     epoch += 1
#     inputs = torch.from_numpy(x_train)
#     labels = torch.from_numpy(y_train)
#     optimizer.zero_grad()
#     # outputs = net(inputs)
#     outputs = net(wb,f)
#     # TODO：loss目前想的是用模型输出的分配资源总和与使用max-min算法的标准答案比较
#     loss = criterion(outputs, labels)
#     loss.backward()
#     optimizer.step()
#     print('epoch {}, loss {}'.format(epoch, loss.item()))